package org.apache.flink.runtime.rest.handler.resourceconsumer;

import java.util.Map;
import java.util.concurrent.CompletableFuture;

import javax.annotation.Nonnull;

import org.apache.flink.api.common.time.Time;
import org.apache.flink.runtime.resourceconsumer.ResourceConsumer;
import org.apache.flink.runtime.resourceconsumer.ResourceConsumerGateway;
import org.apache.flink.runtime.rest.handler.AbstractRestHandler;
import org.apache.flink.runtime.rest.handler.HandlerRequest;
import org.apache.flink.runtime.rest.handler.RestHandlerException;
import org.apache.flink.runtime.rest.messages.MessageHeaders;
import org.apache.flink.runtime.rest.messages.MessageParameters;
import org.apache.flink.runtime.rest.messages.RequestBody;
import org.apache.flink.runtime.rest.messages.ResponseBody;
import org.apache.flink.runtime.rpc.retriever.GatewayRetriever;
import org.apache.flink.runtime.webmonitor.RestfulGateway;
import org.apache.flink.shaded.netty4.io.netty.handler.codec.http.HttpResponseStatus;
import org.apache.flink.util.Preconditions;

/**
 * Base class for REST handlers which need access to the {@link ResourceConsumer}.
 *
 * @param <T> type of the {@link RestfulGateway}
 * @param <R> request type
 * @param <P> response type
 * @param <M> message parameters type
 */
public abstract class AbstractResourceConsumerHandler<
                T extends RestfulGateway,
                R extends RequestBody,
                P extends ResponseBody,
                M extends MessageParameters>
        extends AbstractRestHandler<T, R, P, M> {

    private final GatewayRetriever<ResourceConsumerGateway> resourceConsuemrGatewayRetriever;

    protected AbstractResourceConsumerHandler(
            GatewayRetriever<? extends T> leaderRetriever,
            Time timeout,
            Map<String, String> responseHeaders,
            MessageHeaders<R, P, M> messageHeaders,
            GatewayRetriever<ResourceConsumerGateway> resourceConsuemrGatewayRetriever) {
        super(leaderRetriever, timeout, responseHeaders, messageHeaders);

        this.resourceConsuemrGatewayRetriever =
                Preconditions.checkNotNull(resourceConsuemrGatewayRetriever);
    }

    @Override
    protected CompletableFuture<P> handleRequest(
            @Nonnull HandlerRequest<R> request, @Nonnull T gateway) throws RestHandlerException {
    	ResourceConsumerGateway resourceManagerGateway =
                getResourceConsumerGateway(resourceConsuemrGatewayRetriever);

        return handleRequest(request, resourceManagerGateway);
    }

    protected abstract CompletableFuture<P> handleRequest(
            @Nonnull HandlerRequest<R> request, @Nonnull ResourceConsumerGateway gateway)
            throws RestHandlerException;

    public static ResourceConsumerGateway getResourceConsumerGateway(
            GatewayRetriever<ResourceConsumerGateway> resourceConsuemrGatewayRetriever)
            throws RestHandlerException {
        return resourceConsuemrGatewayRetriever
                .getNow()
                .orElseThrow(
                        () ->
                                new RestHandlerException(
                                        "Cannot connect to ResourceConsumer right now. Please try to refresh.",
                                        HttpResponseStatus.NOT_FOUND));
    }
}
